import numpy as np
from PIL import Image
import os
from sklearn.decomposition import PCA
def img_to_grayscale_vector(image_path, size=(8, 8)):
    """
    将图像转换为指定大小的灰度图，并展平为一维向量。
    """
    img = Image.open(image_path).convert('L') # 转换为灰度图 'L'
    img = img.resize(size, Image.Resampling.LANCZOS)
    return np.array(img).flatten()

class PCAPerceptualHasher:
    def __init__(self, hash_size=4):
        self.hash_size = hash_size
        self.pca = PCA(n_components=hash_size)

    def train(self, image_paths):
        all_image_vectors = []
        for path in image_paths:
            vec = img_to_grayscale_vector(path)
            all_image_vectors.append(vec)
        data_matrix = np.array(all_image_vectors)
        self.pca.fit(data_matrix)

    def calculate_pca_hash(self, image_path):
        vec = img_to_grayscale_vector(image_path)
        transformed = self.pca.transform([vec])[0]
        binary_hash = (transformed > 0).astype(int)
        return ''.join(map(str, binary_hash))
    
if __name__ == "__main__":
    image_files = [
        "demo_images/square.png",
        "demo_images/large_square_gray.png",
        "demo_images/top_left_square.png",
        "demo_images/square_with_noise.png"
    ]

    hasher = PCAPerceptualHasher()
    hasher.train(image_files)

    print("--- 计算 PCA 哈希值 ---")
    for path in image_files:
        img_hash = hasher.calculate_pca_hash(path)
        print(f"{path} PCA 哈希: {img_hash}")